import torch
import joblib
from MLP import CardioMLP
import pandas as pd
import numpy as np

model_state_dict = torch.load('cardio_model.pth')
model = CardioMLP()
model.load_state_dict(model_state_dict)
model.eval()
scaler = joblib.load("scaler.save")


def predict_user_input(model):
    # 提示用户输入数据
    print("请输入以下信息以进行预测：")
    age = float(input("年龄（年）："))
    gender = float(input("性别（1为男性，0为女性）："))
    height = float(input("身高（厘米）："))
    weight = float(input("体重（千克）："))
    ap_hi = float(input("收缩压（mmHg）："))
    ap_lo = float(input("舒张压（mmHg）："))
    cholesterol = float(input("胆固醇水平（1为正常，2为高，3为非常高）："))
    gluc = float(input("血糖水平（1为正常，2为高，3为非常高）："))
    smoke = float(input("是否吸烟（1为是，0为否）："))
    alco = float(input("是否饮酒（1为是，0为否）："))
    active = float(input("是否活跃（1为是，0为否）："))

    bmi = weight / ((height / 100) ** 2)

    feature_names = ['age', 'gender', 'BMI', 'ap_hi', 'ap_lo', 'cholesterol', 'gluc', 'smoke', 'alco', 'active']
    user_input = pd.DataFrame(np.array([[age, gender, bmi, ap_hi, ap_lo, cholesterol, gluc, smoke, alco, active]]),
                              columns=feature_names)

    user_input_normalized = scaler.transform(user_input)

    user_input_tensor = torch.tensor(user_input_normalized, dtype=torch.float).cuda()

    model = model.cuda()

    with torch.no_grad():
        output = model(user_input_tensor)
        probabilities = torch.softmax(output, dim=1)
        prob_0, prob_1 = probabilities.squeeze().tolist()

    print(f"预测结果：患病概率为 {prob_1:.2%}，不患病概率为 {prob_0:.2%}")


predict_user_input(model)
